use datasource_connector::connector::SimpleDataSourceConnector;
use derive_builder::Builder;
use model_graph_types::{
    datasource::{AttributeConditionBuilder, DataSourceConnector, QueryConditionBuilder},
    modeling::{
        assistant::ModelProviderSetting, system_meta_database, ClassGetter, ComparsionOperator,
    },
};
use ollama_rs::Ollama;
use serde::{Deserialize, Serialize};

#[derive(Debug, Default, Serialize, Deserialize, Clone, Builder)]
#[builder(name = "OllamaModelBuilder")]
#[builder(pattern = "mutable")]
#[builder(setter(into, strip_option), default)]
#[builder(derive(Debug))]
#[builder(build_fn(error = "anyhow::Error"))]
pub struct OllamaModel {
    pub name: String,
    pub modified_at: String,
    pub size: u64,
}

#[derive(Debug, Default, Serialize, Deserialize, Clone, Builder)]
#[builder(name = "OllamaModelStatusBuilder")]
#[builder(pattern = "mutable")]
#[builder(setter(into, strip_option), default)]
#[builder(derive(Debug))]
#[builder(build_fn(error = "anyhow::Error"))]
pub struct OllamaModelStatus {}

#[derive(Debug, Default, Serialize, Deserialize, Clone, Builder)]
#[builder(name = "OllamaModelInforBuilder")]
#[builder(pattern = "mutable")]
#[builder(setter(into, strip_option), default)]
#[builder(derive(Debug))]
#[builder(build_fn(error = "anyhow::Error"))]
pub struct OllamaModelInfor {
    #[serde(default = "String::new")]
    pub license: String,
    #[serde(default = "String::new")]
    pub modelfile: String,
    #[serde(default = "String::new")]
    pub parameters: String,
    #[serde(default = "String::new")]
    pub template: String,
}
#[derive(Debug, Default, Serialize, Deserialize, Clone, Builder)]
#[builder(name = "OllamaPullModelStatusBuilder")]
#[builder(pattern = "mutable")]
#[builder(setter(into, strip_option), default)]
#[builder(derive(Debug))]
#[builder(build_fn(error = "anyhow::Error"))]
pub struct OllamaPullModelStatus {
    #[serde(rename = "status")]
    pub message: String,
    pub digest: Option<String>,
    pub total: Option<u64>,
    pub completed: Option<u64>,
}

pub async fn ollama_model_provider() -> anyhow::Result<String> {
    Ok(String::from("ollama"))
}

pub async fn ollama_list_local_models(model_provider: &String) -> anyhow::Result<Vec<OllamaModel>> {
    let ollama = _get_ollama(model_provider).await?;

    let list = ollama
        .list_local_models()
        .await
        .map_err(|err| anyhow::anyhow!("{}", err))?;

    Ok(list
        .iter()
        .filter_map(|item| {
            OllamaModelBuilder::default()
                .name(item.name.clone())
                .modified_at(item.modified_at.clone())
                .size(item.size)
                .build()
                .ok()
        })
        .collect())
}

pub async fn ollama_pull_model(
    model_provider: &String,
    model: &String,
) -> anyhow::Result<OllamaPullModelStatus> {
    let ollama = _get_ollama(model_provider).await?;

    let status = ollama.pull_model(model.clone(), true).await?;

    Ok(OllamaPullModelStatusBuilder::default()
        .message(status.message)
        .digest(status.digest.map_or(String::new(), |v| v))
        .total(status.total.map_or(0, |v| v))
        .completed(status.total.map_or(0, |v| v))
        .build()?)
}

pub async fn ollama_remove_model(model_provider: &String, model: &String) -> anyhow::Result<()> {
    let ollama = _get_ollama(model_provider).await?;
    ollama.delete_model(model.clone()).await?;
    Ok(())
}

pub async fn ollama_show_model(
    model_provider: &String,
    model: &String,
) -> anyhow::Result<OllamaModelInfor> {
    let ollama = _get_ollama(model_provider).await?;
    let res = ollama
        .show_model_info(model.clone())
        .await
        .map_err(|err| anyhow::anyhow!("{}", err))?;
    Ok(OllamaModelInforBuilder::default()
        .license(res.license)
        .modelfile(res.modelfile)
        .parameters(res.parameters)
        .template(res.template)
        .build()?)
}

async fn _get_ollama(model_provider: &String) -> anyhow::Result<Ollama> {
    //
    let sdsc = SimpleDataSourceConnector::default();
    let db = system_meta_database()?;

    let provider: ModelProviderSetting = sdsc
        .instance_query_one(
            &Some(db.clone()),
            &db,
            &ModelProviderSetting::get_identifier()?,
            QueryConditionBuilder::default()
                .attributes(vec![AttributeConditionBuilder::default()
                    .name("provider_identifier")
                    .operator(ComparsionOperator::Equal)
                    .value(model_provider)
                    .build()?])
                .build()?,
            None,
            None,
        )
        .await?
        .try_into()?;

    let config: OllamaConfig = serde_json::from_str(&provider.encrypted_config)
        .map_err(|err| anyhow::anyhow!("{}", err))?;
    //
    let ollama = Ollama::new(config.url_base.clone(), config.port);
    Ok(ollama)
}

#[derive(Debug, Default, Serialize, Deserialize, Clone, Builder)]
#[builder(name = "OllamaConfigBuilder")]
#[builder(pattern = "mutable")]
#[builder(setter(into, strip_option), default)]
#[builder(derive(Debug))]
#[builder(build_fn(error = "anyhow::Error"))]
pub struct OllamaConfig {
    pub url_base: String,
    pub port: u16,
}
